import time

import torch

from util.Queue import infoIn

COLLECTION_NAME = 'image_search_wanwei'  # Collection name
MILVUS_HOST = "10.50.9.18"
MILVUS_PORT = "19530"

TOP_K = 100
from pymilvus import connections, MilvusClient
from pymilvus import Collection
from torchvision import transforms

class MilvusUtil:
    def __init__(self):
        infoIn("连接数据库")
        self.client = MilvusClient(
            uri="http://"+MILVUS_HOST+":" + MILVUS_PORT
        )
        connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
        self.collection = Collection(name=COLLECTION_NAME)
        self.collection.load()

        self.model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
        self.model = torch.nn.Sequential(*(list(self.model.children())[:-1]))
        self.model.eval()

        infoIn('加载transforms')
        self.preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        # Embed the search images
    def embed(self,data):
        with torch.no_grad():
            ret = self.model(torch.stack(data))
            # If more than one image, use squeeze
            if len(ret) > 1:
                return ret.squeeze().tolist()
             # Squeeze would remove batch for single image, so using flatten
            else:
                return torch.flatten(ret, start_dim=1).tolist()


    def __query__(self,image):
        infoIn("search=========================================================")
        t1 = time.time()
        im = image.convert('RGB')
        t2 = time.time()
        infoIn("RGB  " + str((t2 - t1)))
        data = self.preprocess(im);
        t3 = time.time()
        infoIn("transforms  " + str((t3 - t2)))
        embedData = self.embed([data])
        t4 = time.time()
        infoIn("embed  " + str((t4 - t3)))
        res = self.collection.search(embedData, anns_field='image_embedding', param={'nprobe': 128}, limit=TOP_K, output_fields=['url','sku','nwSku'])
        t5 = time.time()
        infoIn("query  " + str((t5 - t4)))
        list = []
        for item in res[0]:
            list.append(item.fields)
        return list
    def __remove__(self,ids):
        infoIn("remove==========================================================")
        res = self.client.delete(COLLECTION_NAME,ids)
        return res
    def __insert__(self,image,url,sku,nwSku):
        infoIn("add==============================================================")
        t1 = time.time()
        im = image.convert('RGB')
        t2 = time.time()
        infoIn("convert  " + str(t2 - t1))
        data = self.preprocess(im);
        t3 = time.time()
        infoIn("preprocess  " + str(t3 - t2))
        embed = self.embed([data])
        t4 = time.time()
        infoIn("embed  " + str(t4 - t3))
        res = self.collection.insert([[url],[sku],[nwSku],embed])
        return res

milvusSingle = MilvusUtil()